import torch
import torch.nn as nn
from matplotlib import pyplot as plt

# 准备数据
# 10组数据，每组数据有一个特征
x = torch.arange(10).float().unsqueeze(1)
y = x * 0.5 + 1
# 目前x和y是标准线性，实际上y的值是有一些抖动的，现在来模拟抖动
y += torch.normal(0, 0.1, y.shape)

# 1. 超参数
lr = 0.01
Epochs = 500

# 2. 创建模型——输入特征数为1，输出特征数为1
model = nn.Linear(1, 1)

# 3. 创建损失函数
loss_fn = nn.MSELoss()

# 4. 创建优化器 (SGD--梯度下降)
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
epoch_list = []
loss_list = []
# 5. 循环训练
for epoch in range(Epochs):
    # 前向传播
    y_pre = model(x)
    # 计算损失
    loss = loss_fn(y, y_pre)
    # 反向传播
    loss.backward()
    # 更新参数
    optimizer.step()
    # 清零梯度
    optimizer.zero_grad()

    epoch_list.append(epoch)
    loss_list.append(loss.item())
    # 打印训练信息
    if epoch == 0 or (epoch + 1) % 10 == 0:
        print(f'[{epoch + 1}/{Epochs}], loss: {loss.item():.4f}')
        ax1.clear()
        ax1.scatter(x.numpy(), y.numpy(), color="b")
        ax1.plot(x.numpy(), y_pre.detach().numpy(), color="r")

        ax2.clear()
        ax2.plot(epoch_list, loss_list)

        plt.pause(0.1)

# 打印训练得到的参数 w≈0.5, b≈1
print(loss_list)
print(f'w: {model.weight.item()}, b: {model.bias.item()}')

# 6. 学生练习，绘制动态拟合图（循环更新） 和 损失曲线图
plt.show()

'''
    在pytorch等框架中必须清空梯度，核心原因是：梯度会默认累加，不清空会导致当前梯度与历史梯度叠加，
    破坏参数更新的正确性。
    1. 梯度的物理意义：loss.backward()计算的梯度，是在当前batch数据下参数对损失的偏导数，仅对应
    当前批次的更新方向。
    2. 框架的默认行为：pytorch不会自动重置梯度，而是将新计算的梯度累加在param.grad。若不清空，下一轮
    梯度会包含上一轮的历史梯度值，导致optimizer.step()偏离预期。

    如果想关闭梯度计算，怎么做呢？
    with torch.no_grad():  是pytorch中非常实用的上下文管理器，主要作用是在其作用域内停止/关闭梯度计算，
    这样做可以显著提升计算效率，节省内存占用。一般主要用于推理阶段，因为这个时候只需要模型的输出，不需要梯度信息的。
    '''